#include "PKG.h"

using CryptoPP::RandomNumberGenerator;
using CryptoPP::BIG_ENDIAN_ORDER;
using CryptoPP::P1363_KDF2;
using CryptoPP::word32;
using CryptoPP::SecByteBlock;

using CryptoPP::ByteQueue;

using CryptoPP::FileSource;
using CryptoPP::FileSink;

using CryptoPP::RSA;

using CryptoPP::Base64Encoder;
using CryptoPP::Base64Decoder;

using CryptoPP::PrivateKey;
using CryptoPP::PublicKey;

using CryptoPP::BufferedTransformation;

using CryptoPP::SHA1;

using CryptoPP::HMAC;

using std::runtime_error;


KDF2_RNG::KDF2_RNG(const byte *seed, size_t seedSize)
                 : m_counter(0), m_counterAndSeed(seedSize + 4)
{
  memcpy(m_counterAndSeed + 4, seed, seedSize);
}
 
void KDF2_RNG::GenerateBlock(byte *output, size_t size)
{
  PutWord(false, BIG_ENDIAN_ORDER, m_counterAndSeed, m_counter);
  ++m_counter;
  P1363_KDF2<SHA1>::DeriveKey(output, size, m_counterAndSeed, m_counterAndSeed.size(), NULL, 0);
}

byte *PKG::IdentifyDigestFunction(byte *ID)
{
  if(!ID)
  {
    return NULL;
  }

  // do stuff
  HMAC<SHA1> mac;
  byte *IDT = (byte*)malloc(mac.DIGESTSIZE);
  //memset(IDT,0,mac.DIGESTSIZE);
  
  mac.SetKey(this->masterKey,strlen((const char*)this->masterKey));
  mac.Update(ID,strlen((const char*)ID));
  mac.Final(IDT);
    
  return IDT;
}
  
PKG::PKG(byte *mKey)
{
  this->masterKey = (byte*)malloc(sizeof(byte)*strlen((const char*)mKey));

  if(!this->masterKey)
  {
    throw runtime_error("Unable to allocate space for the masterKey!");
  }
  
  strcpy((char*)this->masterKey,(const char*)mKey);
}

PubPriPair *PKG::GenerateKeyPair(byte *ID)
{
  byte *IDT = NULL;
  
  IDT = IdentifyDigestFunction(ID);
    
  if(!IDT)
  {
    printf("No IDT generated!\n");
    
    return NULL;
  }
  
  PubPriPair *PPP = new PubPriPair;
  
  if(!PPP)
  {
    throw runtime_error("Unable to allocate space for the masterKey!");
  }
 
  //CryptoPP::LC_RNG rng(strlen((const char*)IDT));
  //RandomPool rng;
  //printf("IDT = %s(%d)\n",IDT,strlen((const char*)IDT));

  KDF2_RNG rng(IDT,strlen((const char*)IDT));
  free(IDT);
  
  // Generate the private key
  RSA::PrivateKey privateKey;
  privateKey.GenerateRandomWithKeySize(rng, 3072);
  //printf("ID:%s\n",ID);
  std::string pristr = std::string("private-").append((const char*)ID).append(".key");
  
  //printf("pre-saved private\n");
  //SavePrivateKey(pristr,privateKey);
  try
  {
    EncodePrivateKey(pristr,privateKey);
  //printf("post-saved private\n");
    
  // Validate the private key
  /*RSA::PrivateKey r2;
  LoadPrivateKey(pristr, r2);

  if(!privateKey.Validate(rng, 2))
  {
    throw runtime_error("Private key does not validate!");
  }
  if(!r2.Validate(rng, 2))
  {
    throw runtime_error("Private key does not validate!");
  }

  if(privateKey.GetModulus() != r2.GetModulus() ||
    privateKey.GetPublicExponent() != r2.GetPublicExponent() ||
    privateKey.GetPrivateExponent() != r2.GetPrivateExponent())
  {
    throw runtime_error("Key data was not stored correctly!");
  }
		
  printf("Private key validated.\n");*/
  
  PPP->privateFile = std::string(pristr);
  // Generate the public key
  RSA::PublicKey publicKey(privateKey);
  std::string pubstr = std::string("public-").append((const char*)ID).append(".key");
  //printf("public key (%s)\n",pubstr.c_str());
  
  EncodePublicKey(pubstr,publicKey);
  //SavePublicKey(pubstr,publicKey);
  
  PPP->publicFile = std::string(pubstr);
  }
  catch(CryptoPP::Exception& e)
  {
    std::cerr << e.what() << std::endl;
    delete PPP;
    return -0;
  }
  
  // debug
  //printf("did public key\n");
    
  return PPP;
}
